// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/vector/string_vector_buffer.h>

namespace kumo::pollux {
    StringVectorBuffer::StringVectorBuffer(
        FlatVector<StringView> *vector,
        size_t initialCapacity,
        size_t maxCapacity)
        : maxCapacity_(maxCapacity), vector_(vector) {
        POLLUX_CHECK_GE(
            maxCapacity_, initialCapacity, "initialCapacity must be <= maxCapacity");
        ensureCapacity(initialCapacity);
    }

    void StringVectorBuffer::appendByte(int8_t value) {
        ensureCapacity(1);
        POLLUX_CHECK_GT(writableCapacity(), 0, "No writable capacity");
        rawBuffer_[currentPosition_++] = value;
    }

    void StringVectorBuffer::flushRow(vector_size_t rowId) {
        // Flush the current row to the vector.
        vector_->setNoCopy(
            rowId, StringView(data() + startPosition_, unflushedRowSize()));
        startPosition_ = currentPosition_;
    }

    void StringVectorBuffer::ensureCapacity(size_t growSize) {
        if (writableCapacity() >= growSize) {
            return;
        }

        // minimal space needed for resize to write the current unflushed row and
        // 'growSize' number of additional bytes.
        const auto currUnflushedRowSize = unflushedRowSize();
        auto minRequiredCapacity = totalCapacity_ + currUnflushedRowSize + growSize;

        auto newTotalCapacity = std::min(
            std::max(totalCapacity_ * kGrowFactor, minRequiredCapacity),
            maxCapacity_);

        POLLUX_CHECK_GE(
            newTotalCapacity,
            minRequiredCapacity,
            "Cannot grow buffer with totalCapacity:{} to meet minRequiredCapacity:{}",
            succinctBytes(totalCapacity_),
            succinctBytes(minRequiredCapacity));

        // newBufferCapacity is the additional free space we need to allocate.
        const auto newBufferCapacity = newTotalCapacity - totalCapacity_;

        auto newBuffer = vector_->getRawStringBufferWithSpace(newBufferCapacity);

        // A row needs to fit in continuous memory.
        // Copy the data from the old buffer to the new buffer,
        // if the current row has not been flushed.
        if (currUnflushedRowSize > 0) {
            POLLUX_CHECK_GE(
                newBufferCapacity,
                currUnflushedRowSize,
                "not enough buffer space to write unflushed row");
            std::memcpy(newBuffer, data() + startPosition_, currUnflushedRowSize);
        }

        rawBuffer_ = newBuffer;
        startPosition_ = 0;
        currentPosition_ = currUnflushedRowSize;

        // Update the capacity.
        currentCapacity_ = newBufferCapacity;
        totalCapacity_ = newTotalCapacity;
    }
} // namespace kumo::pollux
